import os
import random

import tensorflow as tf
import numpy as np

import Layers
import LSTMPs


dataset_path = 'D:/Datasets/TIMIT/stack/'
NUM_CLASSES = 40
NORM_FRAMES = 256
SAMPLE_DEPTH = 39


# augmentation
def _augmentation_online(input_frames):
	max = tf.reduce_max(input_frames)
	min = tf.reduce_min(input_frames)
	input_data = tf.identity(input_frames)
	input_data = tf.subtract(input_data, min)
	input_data = tf.divide(input_data, tf.subtract(max, min))
	input_data = tf.subtract(input_data, 0.5)
	input_data = tf.multiply(input_data, 2.0)
	return input_data


def _sparse_tuple_from(sequences, dtype=np.int32):
    """Create a sparse representention of x.
    Args:
        sequences: a list of lists of type dtype where each element is a sequence
    Returns:
        A tuple with (indices, values, shape)
    """
    indices = []
    values = []

    for n, seq in enumerate(sequences):
        indices.extend(zip([n]*len(seq), range(len(seq))))
        values.extend(seq)

    indices = np.asarray(indices, dtype=np.int64)
    values = np.asarray(values, dtype=dtype)
    shape = np.asarray([len(sequences), np.asarray(indices).max(0)[1]+1], dtype=np.int64)

    return indices, values, shape


def _pad_sequences(sequences, maxlen=None, dtype=np.float32,
                  padding='post', truncating='post', value=0.):
    '''Pads each sequence to the same length: the length of the longest
    sequence.
        If maxlen is provided, any sequence longer than maxlen is truncated to
        maxlen. Truncation happens off either the beginning or the end
        (default) of the sequence. Supports post-padding (default) and
        pre-padding.

        Args:
            sequences: list of lists where each element is a sequence
            maxlen: int, maximum length
            dtype: type to cast the resulting sequence.
            padding: 'pre' or 'post', pad either before or after each sequence.
            truncating: 'pre' or 'post', remove values from sequences larger
            than maxlen either in the beginning or in the end of the sequence
            value: float, value to pad the sequences to the desired value.
        Returns
            x: numpy array with dimensions (number_of_sequences, maxlen)
            lengths: numpy array with the original sequence lengths
    '''
    lengths = np.asarray([len(s) for s in sequences], dtype=np.int64)

    nb_samples = len(sequences)
    if maxlen is None:
        maxlen = np.max(lengths)

    # take the sample shape from the first non empty sequence
    # checking for consistency in the main loop below.
    sample_shape = tuple()
    for s in sequences:
        if len(s) > 0:
            sample_shape = np.asarray(s).shape[1:]
            break

    x = (np.ones((nb_samples, maxlen) + sample_shape) * value).astype(dtype)
    for idx, s in enumerate(sequences):
        if len(s) == 0:
            continue  # empty list was found
        if truncating == 'pre':
            trunc = s[-maxlen:]
        elif truncating == 'post':
            trunc = s[:maxlen]
        else:
            raise ValueError('Truncating type "%s" not understood' % truncating)

        # check `trunc` has expected shape
        trunc = np.asarray(trunc, dtype=dtype)
        if trunc.shape[1:] != sample_shape:
            raise ValueError('Shape of sample %s of sequence at position %s is different from expected shape %s' %
                             (trunc.shape[1:], idx, sample_shape))

        if padding == 'post':
            x[idx, :len(trunc)] = trunc
        elif padding == 'pre':
            x[idx, -len(trunc):] = trunc
        else:
            raise ValueError('Padding type "%s" not understood' % padding)
    return x, lengths


# get dataset
def get_dataset(str_restore_ckpt = None):
	ds_t = np.load(dataset_path + 'train_phn.npz', allow_pickle=True)
	ds_v = np.load(dataset_path + 'test_phn.npz', allow_pickle=True)

	ds_train_fi = ds_t['fi']
	ds_train_label = ds_t['label']
	ds_train_time = ds_t['times']
	ds_val_fi = ds_v['fi']
	ds_val_label = ds_v['label']

	dict_dataset = {}
	dict_dataset['train'] = {
		'train_labels' : ds_train_label,
		'train_data' : ds_train_fi,
		'train_times' : ds_train_time
		}
	dict_dataset['validation'] = {
		'validation_labels' : ds_val_label,
		'validation_data' : ds_val_fi
		}

	return dict_dataset


# orgnize train or val data each time, which is flag_batch_size
def construct_batch_part(flag_batch_size):
	# train/val placeholder
	tfph_data = tf.placeholder(dtype = tf.float32, shape = [flag_batch_size, NORM_FRAMES, SAMPLE_DEPTH], name = 'ph_data')

	# labels placeholders
	tfph_labels = tf.sparse_placeholder(dtype = tf.int32, name = 'ph_labels')
	
	# length placeholders
	tfph_len = tf.placeholder(tf.int32, shape = [flag_batch_size], name = 'ph_len')
	
	# augmentation to [-1,1]
	for k in range(flag_batch_size):
		tfph_frames = tfph_data[k]
		
		# online augmentation
		aug_frames = _augmentation_online(tfph_frames)
		
		# recover to batch data
		if k == 0:
			batch_data = tf.expand_dims(aug_frames, 0)
		else:
			batch_data = tf.concat([batch_data, tf.expand_dims(aug_frames, 0)], axis = 0)
			
	result = {}
	result['batches'] = {
		'batch_data' : batch_data,
		'batch_labels' : tfph_labels,
		'batch_lengths' : tfph_len,
		}
	result['input_placeholders'] = {
		'tfph_data' : tfph_data,
		'tfph_labels' : tfph_labels,
		'tfph_lengths' : tfph_len,
		}
	return result


# orgnize a batch of train data, combining with construct_batch_part
def get_batch_part_train(dict_dataset, dict_placeholders, n_index_head, flag_batch_size):
	n_size = dict_dataset['train']['train_labels'].shape[0]

	n_index_end = n_index_head + flag_batch_size
	if n_index_end > n_size:
		n_index_end = n_size - 1
		n_index_head = n_index_end - flag_batch_size

	lst_data = []
	lst_labels = []
	k = 0
	for data in dict_dataset['train']['train_data'][n_index_head:n_index_end]:
		length = data.shape[0]
		if length > NORM_FRAMES:
			n_anchor = np.random.randint(0, length - NORM_FRAMES)
			lst_data.append(data[n_anchor : n_anchor + NORM_FRAMES])
			time_s = n_anchor * 160
			time_e = (n_anchor + NORM_FRAMES) * 160
			time_pairs = dict_dataset['train']['train_times'][n_index_head + k]
			s = 0
			e = 0
			for i in range(len(time_pairs)):
				if time_pairs[i][0] <= time_s <= time_pairs[i][1]:
					front = time_s - time_pairs[i][0]
					back = time_pairs[i][1] - time_s
					if front <= back:
						s = i
					else:
						s = i + 1
				if time_pairs[i][0] <= time_e <= time_pairs[i][1]:
					front = time_e - time_pairs[i][0]
					back = time_pairs[i][1] - time_e
					if front <= back:
						e = i
					else:
						e = i + 1
			if e <= s:
				e = len(time_pairs) - 1
			if s >= len(time_pairs) - 1:
				s = len(time_pairs) - 1
			if e + 1 - s > NORM_FRAMES:
				e = s + NORM_FRAMES - 1
			lst_labels.append(dict_dataset['train']['train_labels'][n_index_head + k][s : e + 1])
		else:
			lst_data.append(data)
			lst_labels.append(dict_dataset['train']['train_labels'][n_index_head + k])
			
		k = k + 1

	arr_data, arr_len = _pad_sequences(lst_data, NORM_FRAMES)
	arr_labels = _sparse_tuple_from(lst_labels)

	# only train placeholders are used
	dict_feeder = {
		dict_placeholders['tfph_data'] : arr_data,
		dict_placeholders['tfph_labels'] : arr_labels,
		dict_placeholders['tfph_lengths'] : arr_len,
	}

	return dict_feeder


# orgnize a batch of validation data, combining with construct_batch_part
def get_batch_part_validation(dict_dataset, dict_placeholders, n_index_head, flag_batch_size):
	n_size = dict_dataset['validation']['validation_labels'].shape[0]

	n_index_end = n_index_head + flag_batch_size
	if n_index_end > n_size:
		n_index_end = n_size - 1
		n_index_head = n_index_end - flag_batch_size

	arr_data = dict_dataset['validation']['validation_data'][n_index_head:n_index_end]
	arr_data, arr_len = _pad_sequences(arr_data, NORM_FRAMES)
	arr_labels = _sparse_tuple_from(dict_dataset['validation']['validation_labels'][n_index_head:n_index_end])

	# only train placeholders are used
	dict_feeder = {
		dict_placeholders['tfph_data'] : arr_data,
		dict_placeholders['tfph_labels'] : arr_labels,
		dict_placeholders['tfph_lengths'] : arr_len,
	}

	return dict_feeder


# common LSTM
def _network(data, labels, lengths, tfv_train_phase = None):
	name = 'network_normal_LSTM_TIMIT'
	num_hidden = 256
	num_output = 64

	with tf.variable_scope(name, reuse = tf.AUTO_REUSE):
		# input matrix is (num_input * num_hidden), recurrent matrix is (num_output * num_hidden)
		inputs = tf.transpose(data, [0, 2, 1])
		outputs_fw_1 = LSTMPs.lstmp_layer(inputs, num_hidden, num_output, tfv_train_phase, 0.9, reverse = False, name_scope = 'lstm_fw_1')
		outputs_bw_1 = LSTMPs.lstmp_layer(inputs, num_hidden, num_output, tfv_train_phase, 0.9, reverse = True, name_scope = 'lstm_bw_1')
		outputs_1 = tf.concat([outputs_fw_1, outputs_bw_1], axis = 1)
		outputs_fw_2 = LSTMPs.lstmp_layer(outputs_1, num_hidden, num_output, tfv_train_phase, 0.9, reverse = False, name_scope = 'lstm_fw_2')
		outputs_bw_2 = LSTMPs.lstmp_layer(outputs_1, num_hidden, num_output, tfv_train_phase, 0.9, reverse = True, name_scope = 'lstm_bw_2')
		outputs_fw_2  = tf.transpose(outputs_fw_2, [0, 2, 1])
		outputs_bw_2  = tf.transpose(outputs_bw_2, [0, 2, 1])
		outputs = tf.concat([outputs_fw_2, outputs_bw_2], axis = -1)
		outputs = tf.reshape(outputs, [-1, num_output * 2])
		outputs = Layers.fc(outputs, NUM_CLASSES, act_last = False, name_scope = 'fc_out')

	# for CTC loss, seq_len is at axis = 0, while batch_size at axis = 1
	logits = tf.reshape(outputs, [data.shape[0].value, -1, NUM_CLASSES])
	logits = tf.transpose(logits, [1, 0, 2])

	# loss
	ctc = tf.nn.ctc_loss(labels, logits, lengths)
	losses = tf.reduce_mean(ctc, name = 'losses' + name)
	total_loss = tf.add_n([losses], name = 'total_loss' + name)
	loss_averages = tf.train.ExponentialMovingAverage(0.99, name = 'avg_loss' + name)
	tfop_loss_averages = loss_averages.apply([losses] + [total_loss])
	with tf.control_dependencies([tfop_loss_averages]):
		total_loss = tf.identity(total_loss)

	# label error rate, here is PER in actual
	decoded, _ = tf.nn.ctc_beam_search_decoder(logits, lengths)
	ler = tf.edit_distance(tf.cast(decoded[0], tf.int32), labels)

	return total_loss, ler


# TT-LSTM
def _network_tt(data, labels, lengths, tfv_train_phase = None):
	name = 'network_TT_LSTM_TIMIT'
	num_hidden = 256
	num_output = 64

	with tf.variable_scope(name, reuse = tf.AUTO_REUSE):
		# input matrix is (num_input * num_hidden), recurrent matrix is (num_output * num_hidden)
		inputs = tf.transpose(data, [0, 2, 1])
		outputs_fw_1 = LSTMPs.tt_lstmp_layer(inputs, num_hidden, num_output, 
							None, [4,4,4,4], [4,2,2,4], None, [4,4,4], tfv_train_phase, 0.9, reverse = False, name_scope = 'lstm_fw_1')
		outputs_bw_1 = LSTMPs.tt_lstmp_layer(inputs, num_hidden, num_output, 
							None, [4,4,4,4], [4,2,2,4], None, [4,4,4], tfv_train_phase, 0.9, reverse = True, name_scope = 'lstm_bw_1')
		outputs_1 = tf.concat([outputs_fw_1, outputs_bw_1], axis = 1)
		outputs_fw_2 = LSTMPs.tt_lstmp_layer(outputs_1, num_hidden, num_output, 
							None, [4,4,4,4], [4,2,2,4], None, [4,4,4], tfv_train_phase, 0.9, reverse = False, name_scope = 'lstm_fw_2')
		outputs_bw_2 = LSTMPs.tt_lstmp_layer(outputs_1, num_hidden, num_output, 
							None, [4,4,4,4], [4,2,2,4], None, [4,4,4], tfv_train_phase, 0.9, reverse = True, name_scope = 'lstm_bw_2')
		outputs_fw_2  = tf.transpose(outputs_fw_2, [0, 2, 1])
		outputs_bw_2  = tf.transpose(outputs_bw_2, [0, 2, 1])
		outputs = tf.concat([outputs_fw_2, outputs_bw_2], axis = -1)
		outputs = tf.reshape(outputs, [-1, num_output * 2])
		outputs = Layers.fc(outputs, NUM_CLASSES, act_last = False, name_scope = 'fc_out')

	# for CTC loss, seq_len is at axis = 0, while batch_size at axis = 1
	logits = tf.reshape(outputs, [data.shape[0].value, -1, NUM_CLASSES])
	logits = tf.transpose(logits, [1, 0, 2])

	# loss
	ctc = tf.nn.ctc_loss(labels, logits, lengths)
	losses = tf.reduce_mean(ctc, name = 'losses' + name)
	total_loss = tf.add_n([losses], name = 'total_loss' + name)
	loss_averages = tf.train.ExponentialMovingAverage(0.99, name = 'avg_loss' + name)
	tfop_loss_averages = loss_averages.apply([losses] + [total_loss])
	with tf.control_dependencies([tfop_loss_averages]):
		total_loss = tf.identity(total_loss)

	# label error rate, here is PER in actual
	decoded, _ = tf.nn.ctc_beam_search_decoder(logits, lengths)
	ler = tf.edit_distance(tf.cast(decoded[0], tf.int32), labels)

	return total_loss, ler


# KTD-LSTM
def _network_ktd(data, labels, lengths, tfv_train_phase = None):
	name = 'network_KTD_LSTM_TIMIT'
	num_hidden = 256
	num_output = 64

	with tf.variable_scope(name, reuse = tf.AUTO_REUSE):
		# input matrix is (num_input * num_hidden), recurrent matrix is (num_output * num_hidden)
		inputs = tf.transpose(data, [0, 2, 1])
		outputs_fw_1 = LSTMPs.ktd_lstmp_layer(inputs, num_hidden, num_output, 
							None, [4,4,4,4], [4,2,2,4], 4, None, [2 for i in range(4)] + [2 for i in range(4)], tfv_train_phase, 0.9, reverse = False, name_scope = 'lstm_fw_1')
		outputs_bw_1 = LSTMPs.ktd_lstmp_layer(inputs, num_hidden, num_output, 
							None, [4,4,4,4], [4,2,2,4], 4, None, [2 for i in range(4)] + [2 for i in range(4)], tfv_train_phase, 0.9, reverse = True, name_scope = 'lstm_bw_1')
		outputs_1 = tf.concat([outputs_fw_1, outputs_bw_1], axis = 1)
		outputs_fw_2 = LSTMPs.ktd_lstmp_layer(outputs_1, num_hidden, num_output, 
							None, [4,4,4,4], [4,2,2,4], 4, None, [2 for i in range(4)] + [2 for i in range(4)], tfv_train_phase, 0.9, reverse = False, name_scope = 'lstm_fw_2')
		outputs_bw_2 = LSTMPs.ktd_lstmp_layer(outputs_1, num_hidden, num_output, 
							None, [4,4,4,4], [4,2,2,4], 4, None, [2 for i in range(4)] + [2 for i in range(4)], tfv_train_phase, 0.9, reverse = True, name_scope = 'lstm_bw_2')
		outputs_fw_2  = tf.transpose(outputs_fw_2, [0, 2, 1])
		outputs_bw_2  = tf.transpose(outputs_bw_2, [0, 2, 1])
		outputs = tf.concat([outputs_fw_2, outputs_bw_2], axis = -1)
		outputs = tf.reshape(outputs, [-1, num_output * 2])
		outputs = Layers.fc(outputs, NUM_CLASSES, act_last = False, name_scope = 'fc_out')

	# for CTC loss, seq_len is at axis = 0, while batch_size at axis = 1
	logits = tf.reshape(outputs, [data.shape[0].value, -1, NUM_CLASSES])
	logits = tf.transpose(logits, [1, 0, 2])

	# loss
	ctc = tf.nn.ctc_loss(labels, logits, lengths)
	losses = tf.reduce_mean(ctc, name = 'losses' + name)
	total_loss = tf.add_n([losses], name = 'total_loss' + name)
	loss_averages = tf.train.ExponentialMovingAverage(0.99, name = 'avg_loss' + name)
	tfop_loss_averages = loss_averages.apply([losses] + [total_loss])
	with tf.control_dependencies([tfop_loss_averages]):
		total_loss = tf.identity(total_loss)

	# label error rate, here is PER in actual
	decoded, _ = tf.nn.ctc_beam_search_decoder(logits, lengths)
	ler = tf.edit_distance(tf.cast(decoded[0], tf.int32), labels)

	return total_loss, ler


# get the network training and validation output respectively, including loss and evalation
def get_network_output(flag_model, flag_share, tv_data, tv_labels, tv_lens, tfv_train_phase):
	if flag_model == 0:
		loss, ler = _network(tv_data, tv_labels, tv_lens, tfv_train_phase)
	elif flag_model == 1:
		loss, ler = _network_tt(tv_data, tv_labels, tv_lens, tfv_train_phase)
	elif flag_model == 5:
		loss, ler = _network_ktd(tv_data, tv_labels, tv_lens, tfv_train_phase)

	return loss, ler
